# Step 1 Import Modules
import cv2
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import numpy as np
from mediapipe.framework.formats import landmark_pb2
import time
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles
myPose = mp.solutions.pose

BaseOptions = mp.tasks.BaseOptions
PoseLandmarker = mp.tasks.vision.PoseLandmarker
PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
VisionRunningMode = mp.tasks.vision.RunningMode

model_path = "models/pose_landmarker_full.task"
options = PoseLandmarkerOptions(
    base_options=BaseOptions(model_asset_path=model_path),
    running_mode=VisionRunningMode.VIDEO)

video_path = "videos/bdj.mp4"
cap = cv2.VideoCapture(video_path)
timestemp = 0

detector = vision.PoseLandmarker.create_from_options(options)
def process_frame(frame, timestemp):
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
    pose_landmarker_result = detector.detect_for_video(mp_image, timestemp)

    # Step 5: Draw pos landmarks on the input image
    landmarkers = pose_landmarker_result.pose_landmarks
    pose_landmarks = landmark_pb2.NormalizedLandmarkList()
    if len(landmarkers) == 0:
        return  frame

    pose_landmarks.landmark.extend([
        landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z, visibility=landmark.visibility)
        for landmark in landmarkers[0]
    ])
    image = mp_image.numpy_view()
    annotated_image = np.copy(image)
    mp_drawing.draw_landmarks(annotated_image, pose_landmarks, myPose.POSE_CONNECTIONS,
                              landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style())
    return annotated_image

while cap.isOpened():
    success, frame = cap.read()
    if not success:
        print("Ignoring empty camera frame.")
        break

    frame = process_frame(frame, timestemp)
    timestemp = timestemp + 1
    cv2.imshow("Frame", frame)
    key = cv2.waitKey(1) & 0xFF
    if key == ord("q"):
        break

cv2.destroyAllWindows()










